
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import matplotlib.pyplot as plt
import seaborn as sns
db_config = {
    'host': '111.231.14.211',
    'user': 'tushare',
    'password': 'root',
    'database': 'tushare',
    'port': 13307,
    'charset':'utf8mb4'
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000


df = pd.read_sql_query(
    """
    SELECT d.*, m.net_mf_vol , m.sell_elg_vol , m.buy_elg_vol, m.sell_lg_vol, m.buy_lg_vol , i.closes as i_closes, i.vol as i_vol
    FROM date_1 d
    JOIN moneyflows m
    ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
    LFFT JOIN index_daily i ON d.trade_date = i.trade_date and i.ts_code='000001.SH'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code='000001.SZ'
    """,
    conn,
    chunksize=chunk_size
    ) 
df1 = pd.concat(df, ignore_index=True)

print(df1.head())

df1['zd_close'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)
df1['hz_close'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2)
df1['hz_vol'] = round((df1['i_vol'].shift(1)- df1['i_vol'].shift(2)) / df1['i_vol'].shift(2), 2)

df1 = df1.dropna(subset=['zd_close','hz_close','hz_vol'])


numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()

excluded = ['zd_close', 'ts_code', 'trade_date','id','pre_closes','changes','pct_chg','opens','high','low','closes',
            'vol', 'buy_elg_vol', 'i_closes', 'i_vol', 'sell_lg_vol']
predictors = [col for col in numeric_cols if col not in excluded]


formula = 'zd_close ~' + ' + '.join(predictors)
results1 = smf.ols(formula, data=df1).fit()
print(results1.summary())

plt.figure(figsize=(10,6))
plt.scatter(results1.fittedvalues, df1['zd_close'])
plt.title('Fitted Values vs. Observed VaIues')
plt.xlabel('Fitted Values')
plt.ylabel('Observed Values')
plt.show()


features = df1[[ 'amount','net_mf_vol','sell_elg_vol','buy_lg_vol','hz_close','hz_vol']]
correlation_matrix = features.corr()

#Plotting a heatmap for the correlation matrix
plt.figure(figsize=(10,6))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', cbar=True, fmt='.2f', linewidths=.5)
plt.title('Correlation Matrix between Features')
plt.show()